#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
'''
@File    :   priority_redis_query.py    
@Contact :   18574945291@163.com
@License :   (C)Copyright 2018-20xx, Liugroup-NLPR-CASIA

@Modify Time      @Author    @Version    @Desciption
------------      -------    --------    -----------
2020/6/19 15:27   code_wizard      1.0         None
'''

# import lib
import umsgpack
from django_redis import get_redis_connection
from six.moves import queue as BaseQueue



class Priority_query(object):
    Empty = BaseQueue.Empty
    Full = BaseQueue.Full
    max_timeout = 0.3


    def __init__(self, query_name, lazy_limit=True, maxsize=100000):
        self.query_name = query_name
        self.lazy_limit = lazy_limit
        self.last_qsize = 0  # 队列最终长度
        self.maxsize = maxsize  # 初始化队列容量
        self.redis_connect = get_redis_connection()  # 初始化redis链接

    def full(self):
        if self.maxsize and self.queue_size() >= self.maxsize:
            return True
        else:
            return False

    def queue_size(self):
        self.queue_size = self.redis_connect.zcard(self.query_name)  # 获取有序集合的数据量
        return self.queue_size

    def empty(self):
        return True if self.queue_size() == 0 else False

    def put_nowait(self, obj):
        """
        obj是元祖 由权重和值组成
        """
        if self.lazy_limit and self.last_qsize < self.maxsize:
            pass
        elif self.full():
            raise self.Full
        mapping = {umsgpack.packb(obj[1]): float(obj[0])}
        self.last_qsize = self.redis_connect.zadd(self.query_name, dict(mapping))
        return True

    def get_nowait(self):
        """
        -1,-1 取权重值最大的
        0,0 取权重值最小的
        """
        ret = self.redis_connect.zrange(self.query_name, 0, 0)
        self.redis_connect.zrem(self.query_name, ret[0])        # 并删除一个
        if not ret:
            raise self.Empty
        return umsgpack.unpackb(ret[0])
